from numpy import *
import operator


def createDataSet():
    """
    创建数据集，‘numpy.ndarray’类型
    :return: 训练集 训练集对应类别集合
    """
    group = array([
        [1.0, 1.1],
        [1., 1.],
        [0., 0.],
        [0, 0.1]
    ])

    labels = ['A', 'A', 'B', 'B']

    return group, labels


def classify0(inX, dataSet, labels, k):
    """
    :param inX: 输入特征
    :param dataSet: 训练集
    :param labels: 训练集对应类别集合
    :param k: k值
    :return: 输入特征对应类别
    """
    # 返回数据集行数
    dataSetSize = dataSet.shape[0]
    # tile([1,2],(4,2)) 将[1,2]重复4行，重复2列。目的是为了与训练集集体做减法
    diffMat = tile(inX, (dataSetSize, 1)) - dataSet
    # 每个数取平方
    sqDiffMat = diffMat ** 2
    # axis=1：每行求和；axis=0：每列求和。返回一维数组
    sqDistances = sqDiffMat.sum(axis=1)
    # 求平方根即为欧式距离
    distances = sqDistances ** 0.5
    # 返回欧式距离递增的原索引序列
    sortedDistIndices = distances.argsort()
    # 统计前k个类别里各类数量
    classCount = {}
    for i in range(k):
        # 先获取distances中最小值的索引，再到distances获取该值
        voteIlabel = labels[sortedDistIndices[i]]
        # 相应类出现次数+1
        classCount[voteIlabel] = classCount.get(voteIlabel, 0) + 1
    # 对统计好的前k个类排序,返回排好序的元组列表
    sortedclassCount = sorted(classCount.items(), key=operator.itemgetter(1), reverse=True)
    return sortedclassCount[0][0]


if __name__ == "__main__":
    group, labels = createDataSet()
    resClass = classify0([0, 0], group, labels, k=3)
    print(resClass)
